#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Author : gao
# Time : 2020/7/6 22:48

"""
    文件说明：使用决策树进行分类
"""

import CreateTree
import numpy as np
import Shannon as sn

def classifyByTree(treeDict :dict,labels :list,x :list ):
    '''
    使用决策树进行分类
    :param treeDict: 决策树字典
    :param labels:  数据标签
    :param x: 输入数据
    :return:
    '''
    #获取字典的第一个key
    firstr =next(iter(treeDict),None)
    if firstr is None:
        return CreateTree.getMajorityClass(sn.createDataSet())
    #获取firststr的子节点
    nextDict = treeDict[firstr]
    #当前特征对应的下标
    index = labels.index(firstr)
    #返回值（预测的结果）
    result =None
    for value in nextDict.keys():
        if x[index]==value:
            if type(nextDict[value]).__name__ =='dict':
                result = classifyByTree(nextDict[value],labels,x)
            else:
                result = nextDict[value]
    return result


if __name__ == '__main__':
    dataSet, labels = sn.createDataSet()
    labelsCopy = labels[:]
    featLabels = []
    myTree = CreateTree.createDecisionTree(dataSet, labels, featLabels)
    testVec = [9,0,1,1]#测试数据
    result = classifyByTree(myTree, labelsCopy, testVec)
    if result == 'yes':
        print('放贷')
    if result == 'no':
        print('不放贷')